from collections import deque
from datetime import datetime
from itertools import chain
from typing import List, Dict, Union

import numpy as np
import onnx
from onnx import helper, GraphProto, ModelProto

from auto_opt_tool.interface.base_graph import BaseGraph
from auto_opt_tool.onnx.node import PlaceHolder, Initializer, Node

class OnnxGraph(BaseGraph):
    def __init__(
        self,
        nodes: List[Node] = [],
        inputs: List[PlaceHolder] = [],
        outputs: List[PlaceHolder] = [],
        initializers: List[Initializer] = [],
        value_infos: List[PlaceHolder] = [],
        name: str = None,
        **kwargs: Dict[str, object]
    ):
        super().__init__()
        self._nodes = nodes
        self._inputs = inputs
        self._outputs = outputs
        self._initializers = initializers
        self._value_infos = value_infos
        self.name = name

        self._node_map = {}
        self._prev_map = {}
        self._next_map = {}

        self._meta = {'ir_version': kwargs.get('ir_version', 4),
                    'producer_name': kwargs.get('producer_name', 'AutoOptimizer'),
                    'producer_version': kwargs.get('producer_version', 'beta'),
                    'domain': kwargs.get('domain', ''),
                    'model_version': kwargs.get('model_version', 0),
                    'opset_imports': kwargs.get('opset_imports', None)
        }

        self.update_map()

    @classmethod
    def parse(cls, path_or_bytes: Union[str, ModelProto, GraphProto]):
        if isinstance(path_or_bytes, str):
            onnx_model = onnx.load(path_or_bytes)
        if isinstance(path_or_bytes, ModelProto):
            onnx_model = path_or_bytes
        if isinstance(path_or_bytes, GraphProto):
            onnx_graph = path_or_bytes
            meta = {}
        else:
            onnx_graph = onnx_model.graph
            meta = {'ir_version': onnx_model.ir_version,
                    'domain': onnx_model.domain,
                    'model_version': onnx_model.model_version,
                    'doc_string': onnx_model.doc_string,
                    'opset_imports': onnx_model.opset_import
            }

        inputs = [PlaceHolder.parse(i) for i in onnx_graph.input]
        outputs = [PlaceHolder.parse(o) for o in onnx_graph.output]
        initializers = [Initializer.parse(i) for i in onnx_graph.initializer]
        value_infos = [PlaceHolder.parse(v) for v in onnx_graph.value_info]

        nodes = []
        for node in onnx_graph.node:
            nodes.append(Node.parse(node))

        graph = cls(nodes, inputs, outputs, initializers, value_infos, onnx_graph.name, **meta)
        return graph


    def update_map(self):
        # clear map first
        self._node_map = {}
        self._prev_map = {}
        self._next_map = {}

        for n in chain(self._inputs, self._outputs, self._nodes, self._initializers, self._value_infos):
            self._node_map[n.name] = n
        # update prev node info
        for n in self._nodes:
            for o in n.outputs:
                # if output name not in map
                if not self._prev_map.get(o):
                    self._prev_map[o] = n
                else:
                    # TODO: ERROR: duplicate output names
                    pass
        # update next node info
        for n in self._nodes:
            for i in n.inputs:
                if not self._next_map.get(i):
                    self._next_map[i] = [n]
                else:
                    self._next_map[i].append(n)

    def add_placeholder(self, name: str, dtype: str, shape: List[int], ph_type: str = 'input') -> PlaceHolder:
        # TODO: ERROR: dtype is illegal
        dtype = np.dtype(dtype)
        placeholder = PlaceHolder(name, dtype, shape)
        self._node_map[name] = placeholder
        if ph_type == 'input':
            self._inputs.append(placeholder)
        else:
            self._outputs.append(placeholder)
        return placeholder
    
    def add_initializer(self, name: str, value: np.ndarray) -> Initializer:
        initializer = Initializer(name, value)
        self._node_map[name] = initializer
        self._initializers.append(initializer)
        return initializer

    def add_node(self, name: str, op_type: str, attrs: Dict[str, object] = {}, domain:str = None) -> Node:
        node = Node(name, op_type, attrs=attrs, domain=domain)
        self._node_map[name] = node
        self._nodes.append(node)
        return node

    def insert_node(self, refer_node_name: str, insert_node: Union[Node, PlaceHolder, Initializer], refer_io_index: int = 0, mode: str = 'after'):
        # TODO: exception: name not exists in graph
        refer_node = self._node_map[refer_node_name]
        if refer_node.op_type == 'PlaceHolder':
            raise RuntimeError(
                'Please use another mode with appropriate reference node or other insert methods.')
        
        if len(insert_node.inputs) > 1 or len(insert_node.outputs) > 1:
            raise RuntimeError(
                'Only support inserting node with single input and output.')

        if mode == 'after':
            refer_out_name = refer_node.outputs[refer_io_index]
            new_out_name = f'{refer_node.name}/{insert_node.name}'
            # connect insert node
            refer_node.outputs[refer_io_index] = new_out_name
            insert_node.inputs = [new_out_name]
            insert_node.outputs = [refer_out_name]
            # update prev and next map for new output of reference node
            self._prev_map[new_out_name] = refer_node
            self._next_map[new_out_name] = [insert_node]
            # update prev map for original output of reference node
            self._prev_map[refer_out_name] = insert_node
        elif mode == 'before':
            refer_in_name = refer_node.inputs[refer_io_index]
            new_in_name = f'{insert_node.name}/{refer_node.name}'
            # connect insert node
            refer_node.inputs[refer_io_index] = new_in_name
            insert_node.inputs = [refer_in_name]
            insert_node.outputs = [new_in_name]
            # update prev and next map for new input of reference node
            self._prev_map[new_in_name] = insert_node
            self._next_map[new_in_name] = [refer_node]
            # update next map for original input of reference node
            self._next_map[refer_in_name].append(insert_node)
            self._next_map[refer_in_name].remove(refer_node)            
        else:
            raise ValueError(
                f'The value for mode argument should be "after" or "before", but got "{mode}"')
        
        self._node_map[insert_node.name] = insert_node

    def get_nodes(self, op_type: str) -> Union[List[Node], List[PlaceHolder], List[Initializer]]:
        nodes = []
        for node in self._node_map.values():
            if node.op_type == op_type:
                nodes.append(node)
        return nodes

    def remove(self, name: str, maps: Dict[int, int] = {0:0}) -> bool:
        # TODO: exception: name not exist in graph
        node = self._node_map[name]
        self._node_map.pop(name, None)
        if node in self._inputs:
            self._inputs.remove(node)
            self._next_map.pop(name, None)
            return True
        if node in self._outputs:
            self._outputs.remove(node)
            self._prev_map.pop(name, None)
            return True
        if node in self._initializers:
            self._initializers.remove(node)
            self._next_map.pop(name, None)
            return True
        if node in self._nodes:
            self._nodes.remove(node)
            for in_id, in_name in enumerate(node.inputs):
                # update next map, node is no longer a next node
                self._next_map[in_name].remove(node)
                out_id = maps.get(in_id, None)
                # out_id exists, do connection
                if out_id is not None:
                    out_name = node.outputs[out_id]
                    for next_node in self.get_next_nodes(out_name):
                        next_node_in_id = next_node.get_input_id(out_name)
                        next_node.inputs[next_node_in_id] = in_name
                        # update next map, prev node has new next node
                        self._next_map[in_name].append(next_node)
            # update prev and next map, outputs of node no long exist
            for out_name in node.outputs:
                self._prev_map.pop(out_name, None)
                self._next_map.pop(out_name, None)
            return True
        return False
                    

    def __getitem__(self, key: str) -> Union[Node, PlaceHolder, Initializer]:
        return self._node_map[key]

    def __setitem__(self, key: str, value: Union[Node, PlaceHolder, Initializer]):
        # TODO
        pass

    @property
    def inputs(self) -> List[PlaceHolder]:
        return self._inputs

    @property
    def outputs(self) -> List[PlaceHolder]:
        return self._outputs

    @property
    def nodes(self) -> List[Node]:
        return self._nodes

    @property
    def initializers(self) -> List[Initializer]:
        return self._initializers

    def to_graph(self) -> GraphProto:
        self.toposort()
        return helper.make_graph(nodes=[node.to_proto() for node in self._nodes],
                                name=self.name,
                                inputs=[input.to_proto() for input in self._inputs],
                                outputs=[output.to_proto() for output in self._outputs],
                                initializer=[ini.to_proto() for ini in self._initializers]
        )

    def to_model(self) -> ModelProto:
        return helper.make_model(self.to_graph(), **self._meta)

    def save(self, path: str):
        onnx.save(self.to_model(), path)

    def extract(self, input_name_list: List[str], output_name_list: List[str], new_model_save_path: str, enable_model_check: bool = True):
        # TODO: check input_name_list and output_name_list
        
        start_time = datetime.now()

        # get possible top_nodes and bottom_nodes
        top_nodes = []
        bottom_nodes = []
        for input_name in input_name_list:
            if self.get_next_nodes(input_name):
                top_nodes.extend(self.get_next_nodes(input_name)) 
            else:
                # TODO: illegal input
                pass
        for output_name in output_name_list:
            if self.get_prev_node(output_name):
                bottom_nodes.append(self.get_prev_node(output_name)) 
            else:
                # TODO: illegal output
                pass

        # top-down search
        queue = deque(top_nodes)
        top_down_visited = set()
        while queue:
            node = queue.popleft()
            if node in top_down_visited:
                continue
            top_down_visited.add(node)
            for output_name in node.outputs:
                for next_node in self.get_next_nodes(output_name):
                    queue.append(next_node)

        # bottom-up search
        queue = deque(bottom_nodes)
        bottom_up_visited = set()
        while queue:
            node = queue.popleft()
            if node in bottom_up_visited:
                continue
            bottom_up_visited.add(node)
            for input_name in node.inputs:
                prev_node = self.get_prev_node(input_name)
                if prev_node:
                    queue.append(prev_node) 
    
        # collect reachable nodes
        reachable_nodes = top_down_visited & bottom_up_visited
        self._nodes = [self._node_map[n.name] for n in reachable_nodes]
        
        # collect reachable initializers and value_infos
        self._initializers = []
        self._value_infos = []
        for node in self._nodes:
            for input in node.inputs:
                if self._node_map.get(input, None):
                    n = self._node_map[input]
                    if n.op_type == 'Initializer':
                        self._initializers.append(n)
                    elif n.op_type == 'PlaceHolder':
                        self._value_infos.append(n)
       
        # add inputs and outputs for extract graph
        self._inputs = []
        self._outputs = []
        for input_name in input_name_list:
            if self._node_map.get(input_name, None):
                n = self._node_map[input_name]
                if n.op_type == 'PlaceHolder':
                    self._inputs.append(PlaceHolder(n.name, n.dtype, n.shape))
            else:
                self._inputs.append(PlaceHolder(input_name, np.dtype('float32')))
        for output_name in output_name_list:
            if self._node_map.get(output_name, None):
                n = self._node_map[output_name]
                if n.op_type == 'PlaceHolder':
                    self._outputs.append(PlaceHolder(n.name, n.dtype, n.shape))
            else:
                self._outputs.append(PlaceHolder(output_name, np.dtype('float32')))
        
        # save model
        self.name = 'extract_graph'
        self.update_map()
        self.save(new_model_save_path)
        end_time = datetime.now()
        cost_time = (end_time - start_time).seconds
        print(f'Extract model completed. Cost {cost_time} seconds.')

    def get_prev_node(self, input_name: str) -> Union[Node, PlaceHolder, Initializer]:
        # TODO: raise exception
        return self._prev_map.get(input_name, None)

    def get_next_nodes(self, output_name: str) -> Union[List[Node], List[PlaceHolder], List[Initializer]]:
        # TODO: raise exception
        return self._next_map.get(output_name, [])

    def toposort(self):
        def visited_all_prev_nodes(node, visited):
            for input_name in node.inputs:
                prev_node = self.get_prev_node(input_name)
                if not visited.get(prev_node, False) and prev_node:
                    return False
            return True

        queue = []
        visited = {}
        for node in self._nodes:
            if visited_all_prev_nodes(node, visited):
                queue.append(node)
        
        sorted_nodes = []
        while queue:
            node = queue.pop(0)
            if visited_all_prev_nodes(node, visited):
                sorted_nodes.append(node)
                visited[node] = True
                for output_name in node.outputs:
                    for next_node in self.get_next_nodes(output_name):
                        if next_node not in queue and not visited.get(next_node):
                            queue.append(next_node)
            else:
                queue.append(node)
        
        self._nodes = sorted_nodes

    def infershape(self):
        model = self.to_model()
        # TODO: exception
        inferred_model = onnx.shape_inference.infer_shapes(model)
        graph = inferred_model.graph
        self._value_infos = [PlaceHolder.parse(v) for v in graph.value_info]
        for n in self._value_infos:
            self._node_map[n.name] = n

    def run(self):
        pass

    def dump(self):
        pass 